from langchain_community.llms.tongyi import Tongyi
from langchain_core.prompts.prompt import PromptTemplate
from langchain.chains import LLMChain, MultiPromptChain

import os

from chains.MultiPromptChain import default_chain

os.environ["DASHSCOPE_API_KEY"] = "sk-9d8f1914800e497f8717144e860f99bc"
# 定义模型
llm = Tongyi()
# 定义提示词模板
physics_template = """你是一位非常聪明的物理学教授。 \
您擅长以简洁易懂的方式回答有关物理的问题。 \
当你不知道某个问题的答案时，你就承认你不知道。

这是一个问题:
{input}"""

math_template = """你是一位非常优秀的数学家。 你很擅长回答数学问题。 \
你太棒了，因为你能够将难题分解成各个组成部分，\
回答各个组成部分，然后将它们组合起来回答更广泛的问题。

这是一个问题:
{input}"""

prompt_infos = [
    {
        "name": "physics",
        "description": "适合回答物理问题",
        "prompt_template": physics_template,
    },
    {
        "name": "math",
        "description": "适合回答数学问题",
        "prompt_template": math_template,
    },
]
# 构建两个链
destination_chains = {}
for p_info in prompt_infos:
    name = p_info["name"]
    prompt_template = p_info["prompt_template"]
    prompt = PromptTemplate(template=prompt_template, input_variables=['input'])
    chain = LLMChain(llm=llm, prompt=prompt)
    destination_chains[name] = chain
# 组装routerchain

from langchain.chains.router.llm_router import LLMRouterChain, RouterOutputParser
from langchain.chains.router.multi_prompt_prompt import MULTI_PROMPT_ROUTER_TEMPLATE

# 提取2个链的参数描述
destinations = [f"{p['name']}: {p['description']}" for p in prompt_infos]  # 列表字典
destinations_str = "\n".join(destinations)  # 字符串
print(destinations_str)

# 构建多模板
router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(destinations=destinations_str)
router_prompt = PromptTemplate(
    template=router_template,
    input_variables=["input"],
    output_parser=RouterOutputParser(),
)
# 根据每个chain中的模板描述选择合适的chain
router_chain = LLMRouterChain.from_llm(llm, router_prompt)

chain = MultiPromptChain(
    router_chain=router_chain,
    destination_chains=destination_chains,
    default_chain=default_chain,
)
result = chain.invoke("什么是高等数学")
print(result)
